// Copyright (C) 2024 Kumo inc.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <ktest/ktest.h>
#include <atomic>
#include <turbo/times/time.h>
#include <turbo/base/macros.h>
#include <turbo/log/logging.h>
#include <kthread/internal/kutex.h>
#include <kthread/internal/task_control.h>
#include <kthread/internal/task_group.h>
#include <kthread/internal/kthread.h>
#include <kthread/internal/unstable.h>
#include "kthread/internal/interrupt_pthread.h"

namespace kthread {
    extern std::atomic<TaskControl *> g_task_control;

    inline TaskControl *get_task_control() {
        return g_task_control.load(std::memory_order_consume);
    }
} // namespace kthread

namespace {
    TEST(ButexTest, wait_on_already_timedout_kutex) {
        uint32_t *kutex = kthread::kutex_create_checked<uint32_t>();
        ASSERT_TRUE(kutex);
        timespec now;
        ASSERT_EQ(0, clock_gettime(CLOCK_REALTIME, &now));
        *kutex = 1;
        ASSERT_EQ(-1, kthread::kutex_wait(kutex, 1, &now));
        ASSERT_EQ(ETIMEDOUT, errno);
    }

    void *sleeper(void *arg) {
        kthread_usleep((uint64_t) arg);
        return nullptr;
    }

    void *joiner(void *arg) {
        const long t1 = turbo::Time::current_microseconds();
        for (kthread_t *th = (kthread_t *) arg; *th; ++th) {
            if (0 != kthread_join(*th, nullptr)) {
                KLOG(FATAL) << "fail to join thread_" << th - (kthread_t *) arg;
            }
            long elp = turbo::Time::current_microseconds() - t1;
            EXPECT_LE(labs(elp - (th - (kthread_t *) arg + 1) * 100000L), 15000L)
                                << "timeout when joining thread_" << th - (kthread_t *) arg;
            KLOG(INFO) << "Joined thread " << *th << " at " << elp << "us ["
                      << kthread_self() << "]";
        }
        for (kthread_t *th = (kthread_t *) arg; *th; ++th) {
            EXPECT_EQ(0, kthread_join(*th, nullptr));
        }
        return nullptr;
    }

    struct A {
        uint64_t a;
        char dummy[0];
    };

    struct B {
        uint64_t a;
    };


    TEST(ButexTest, with_or_without_array_zero) {
        ASSERT_EQ(sizeof(B), sizeof(A));
    }


    TEST(ButexTest, join) {
        const size_t N = 6;
        const size_t M = 6;
        kthread_t th[N + 1];
        kthread_t jth[M];
        pthread_t pth[M];
        for (size_t i = 0; i < N; ++i) {
            kthread_attr_t attr = (i == 0 ? KTHREAD_ATTR_PTHREAD : KTHREAD_ATTR_NORMAL);
            ASSERT_EQ(0, kthread_start_urgent(
                    &th[i], &attr, sleeper,
                    (void *) (100000L/*100ms*/ * (i + 1))));
        }
        th[N] = 0;  // joiner will join tids in `th' until seeing 0.
        for (size_t i = 0; i < M; ++i) {
            ASSERT_EQ(0, kthread_start_urgent(&jth[i], nullptr, joiner, th));
        }
        for (size_t i = 0; i < M; ++i) {
            ASSERT_EQ(0, pthread_create(&pth[i], nullptr, joiner, th));
        }

        for (size_t i = 0; i < M; ++i) {
            ASSERT_EQ(0, kthread_join(jth[i], nullptr))
                                        << "i=" << i << " error=" << km_error();
        }
        for (size_t i = 0; i < M; ++i) {
            ASSERT_EQ(0, pthread_join(pth[i], nullptr));
        }
    }


    struct WaiterArg {
        int expected_result;
        int expected_value;
        std::atomic<int> *kutex;
        const timespec *ptimeout;
    };

    void *waiter(void *arg) {
        WaiterArg *wa = (WaiterArg *) arg;
        const long t1 = turbo::Time::current_microseconds();
        const int rc = kthread::kutex_wait(
                wa->kutex, wa->expected_value, wa->ptimeout);
        const long t2 = turbo::Time::current_microseconds();
        if (rc == 0) {
            EXPECT_EQ(wa->expected_result, 0) << kthread_self();
        } else {
            EXPECT_EQ(wa->expected_result, errno) << kthread_self();
        }
        KLOG(INFO) << "after wait, time=" << (t2 - t1) << "us";
        return nullptr;
    }

    TEST(ButexTest, sanity) {
        const size_t N = 5;
        WaiterArg args[N * 4];
        pthread_t t1, t2;
        std::atomic<int> *b1 =
                kthread::kutex_create_checked<std::atomic<int> >();
        ASSERT_TRUE(b1);
        kthread::kutex_destroy(b1);

        b1 = kthread::kutex_create_checked<std::atomic<int> >();
        *b1 = 1;
        ASSERT_EQ(0, kthread::kutex_wake(b1));

        WaiterArg *unmatched_arg = new WaiterArg;
        unmatched_arg->expected_value = *b1 + 1;
        unmatched_arg->expected_result = EWOULDBLOCK;
        unmatched_arg->kutex = b1;
        unmatched_arg->ptimeout = nullptr;
        pthread_create(&t2, nullptr, waiter, unmatched_arg);
        kthread_t th;
        ASSERT_EQ(0, kthread_start_urgent(&th, nullptr, waiter, unmatched_arg));

        const timespec abstime = turbo::Time::to_timespec(turbo::Time::future_time(turbo::Duration::seconds(1)));
        for (size_t i = 0; i < 4 * N; ++i) {
            args[i].expected_value = *b1;
            args[i].kutex = b1;
            if ((i % 2) == 0) {
                args[i].expected_result = 0;
                args[i].ptimeout = nullptr;
            } else {
                args[i].expected_result = ETIMEDOUT;
                args[i].ptimeout = &abstime;
            }
            if (i < 2 * N) {
                pthread_create(&t1, nullptr, waiter, &args[i]);
            } else {
                ASSERT_EQ(0, kthread_start_urgent(&th, nullptr, waiter, &args[i]));
            }
        }

        sleep(2);
        for (size_t i = 0; i < 2 * N; ++i) {
            ASSERT_EQ(1, kthread::kutex_wake(b1));
        }
        ASSERT_EQ(0, kthread::kutex_wake(b1));
        sleep(1);
        kthread::kutex_destroy(b1);
    }


    struct ButexWaitArg {
        int *kutex;
        int expected_val;
        long wait_msec;
        int error_code;
    };

    void *wait_kutex(void *void_arg) {
        ButexWaitArg *arg = static_cast<ButexWaitArg *>(void_arg);
        const timespec ts = turbo::Time::to_timespec(turbo::Time::future_time(turbo::Duration::milliseconds(arg->wait_msec)));
        int rc = kthread::kutex_wait(arg->kutex, arg->expected_val, &ts);
        int saved_errno = errno;
        if (arg->error_code) {
            EXPECT_EQ(-1, rc);
            EXPECT_EQ(arg->error_code, saved_errno);
        } else {
            EXPECT_EQ(0, rc);
        }
        return nullptr;
    }

    TEST(ButexTest, wait_without_stop) {
        int *kutex = kthread::kutex_create_checked<int>();
        *kutex = 7;
        turbo::TimeCost tm;
        const long WAIT_MSEC = 500;
        for (int i = 0; i < 2; ++i) {
            const kthread_attr_t attr =
                    (i == 0 ? KTHREAD_ATTR_PTHREAD : KTHREAD_ATTR_NORMAL);
            ButexWaitArg arg = {kutex, *kutex, WAIT_MSEC, ETIMEDOUT};
            kthread_t th;

            tm.reset();
            ASSERT_EQ(0, kthread_start_urgent(&th, &attr, wait_kutex, &arg));
            ASSERT_EQ(0, kthread_join(th, nullptr));
            tm.stop();

            ASSERT_LT(labs(tm.m_elapsed() - WAIT_MSEC), 250);
        }
        kthread::kutex_destroy(kutex);
    }

    TEST(ButexTest, stop_after_running) {
        int *kutex = kthread::kutex_create_checked<int>();
        *kutex = 7;
        turbo::TimeCost tm;
        const long WAIT_MSEC = 500;
        const long SLEEP_MSEC = 10;
        for (int i = 0; i < 2; ++i) {
            const kthread_attr_t attr =
                    (i == 0 ? KTHREAD_ATTR_PTHREAD : KTHREAD_ATTR_NORMAL);
            kthread_t th;
            ButexWaitArg arg = {kutex, *kutex, WAIT_MSEC, EINTR};

            tm.reset();
            ASSERT_EQ(0, kthread_start_urgent(&th, &attr, wait_kutex, &arg));
            ASSERT_EQ(0, kthread_usleep(SLEEP_MSEC * 1000L));
            ASSERT_EQ(0, kthread_stop(th));
            ASSERT_EQ(0, kthread_join(th, nullptr));
            tm.stop();

            ASSERT_LT(labs(tm.m_elapsed() - SLEEP_MSEC), 25);
            // ASSERT_TRUE(kthread::get_task_control()->
            //             timer_thread()._idset.empty());
            ASSERT_EQ(EINVAL, kthread_stop(th));
        }
        kthread::kutex_destroy(kutex);
    }

    TEST(ButexTest, stop_before_running) {
        int *kutex = kthread::kutex_create_checked<int>();
        *kutex = 7;
        turbo::TimeCost tm;
        const long WAIT_MSEC = 500;

        for (int i = 0; i < 2; ++i) {
            const kthread_attr_t attr =
                    (i == 0 ? KTHREAD_ATTR_PTHREAD : KTHREAD_ATTR_NORMAL) | KTHREAD_NOSIGNAL;
            kthread_t th;
            ButexWaitArg arg = {kutex, *kutex, WAIT_MSEC, EINTR};

            tm.reset();
            ASSERT_EQ(0, kthread_start_background(&th, &attr, wait_kutex, &arg));
            ASSERT_EQ(0, kthread_stop(th));
            kthread_flush();
            ASSERT_EQ(0, kthread_join(th, nullptr));
            tm.stop();

            ASSERT_LT(tm.m_elapsed(), 5);
            // ASSERT_TRUE(kthread::get_task_control()->
            //             timer_thread()._idset.empty());
            ASSERT_EQ(EINVAL, kthread_stop(th));
        }
        kthread::kutex_destroy(kutex);
    }

    void *join_the_waiter(void *arg) {
        EXPECT_EQ(0, kthread_join((kthread_t) arg, nullptr));
        return nullptr;
    }

    TEST(ButexTest, join_cant_be_wakeup) {
        const long WAIT_MSEC = 100;
        int *kutex = kthread::kutex_create_checked<int>();
        *kutex = 7;
        turbo::TimeCost tm;
        ButexWaitArg arg = {kutex, *kutex, 1000, EINTR};

        for (int i = 0; i < 2; ++i) {
            const kthread_attr_t attr =
                    (i == 0 ? KTHREAD_ATTR_PTHREAD : KTHREAD_ATTR_NORMAL);
            tm.reset();
            kthread_t th, th2;
            ASSERT_EQ(0, kthread_start_urgent(&th, nullptr, wait_kutex, &arg));
            ASSERT_EQ(0, kthread_start_urgent(&th2, &attr, join_the_waiter, (void *) th));
            ASSERT_EQ(0, kthread_stop(th2));
            ASSERT_EQ(0, kthread_usleep(WAIT_MSEC / 2 * 1000L));
            ASSERT_TRUE(kthread::TaskGroup::exists(th));
            ASSERT_TRUE(kthread::TaskGroup::exists(th2));
            ASSERT_EQ(0, kthread_usleep(WAIT_MSEC / 2 * 1000L));
            ASSERT_EQ(0, kthread_stop(th));
            ASSERT_EQ(0, kthread_join(th2, nullptr));
            ASSERT_EQ(0, kthread_join(th, nullptr));
            tm.stop();
            ASSERT_LT(tm.m_elapsed(), WAIT_MSEC + 15);
            ASSERT_EQ(EINVAL, kthread_stop(th));
            ASSERT_EQ(EINVAL, kthread_stop(th2));
        }
        kthread::kutex_destroy(kutex);
    }

    TEST(ButexTest, stop_after_slept) {
        turbo::TimeCost tm;
        const long SLEEP_MSEC = 100;
        const long WAIT_MSEC = 10;

        for (int i = 0; i < 2; ++i) {
            const kthread_attr_t attr =
                    (i == 0 ? KTHREAD_ATTR_PTHREAD : KTHREAD_ATTR_NORMAL);
            tm.reset();
            kthread_t th;
            ASSERT_EQ(0, kthread_start_urgent(
                    &th, &attr, sleeper, (void *) (SLEEP_MSEC * 1000L)));
            ASSERT_EQ(0, kthread_usleep(WAIT_MSEC * 1000L));
            ASSERT_EQ(0, kthread_stop(th));
            ASSERT_EQ(0, kthread_join(th, nullptr));
            tm.stop();
            if (attr.stack_type == KTHREAD_STACKTYPE_PTHREAD) {
                ASSERT_LT(labs(tm.m_elapsed() - SLEEP_MSEC), 15);
            } else {
                ASSERT_LT(labs(tm.m_elapsed() - WAIT_MSEC), 15);
            }
            // ASSERT_TRUE(kthread::get_task_control()->
            //             timer_thread()._idset.empty());
            ASSERT_EQ(EINVAL, kthread_stop(th));
        }
    }

    TEST(ButexTest, stop_just_when_sleeping) {
        turbo::TimeCost tm;
        const long SLEEP_MSEC = 100;

        for (int i = 0; i < 2; ++i) {
            const kthread_attr_t attr =
                    (i == 0 ? KTHREAD_ATTR_PTHREAD : KTHREAD_ATTR_NORMAL);
            tm.reset();
            kthread_t th;
            ASSERT_EQ(0, kthread_start_urgent(
                    &th, &attr, sleeper, (void *) (SLEEP_MSEC * 1000L)));
            ASSERT_EQ(0, kthread_stop(th));
            ASSERT_EQ(0, kthread_join(th, nullptr));
            tm.stop();
            if (attr.stack_type == KTHREAD_STACKTYPE_PTHREAD) {
                ASSERT_LT(labs(tm.m_elapsed() - SLEEP_MSEC), 15);
            } else {
                ASSERT_LT(tm.m_elapsed(), 15);
            }
            // ASSERT_TRUE(kthread::get_task_control()->
            //             timer_thread()._idset.empty());
            ASSERT_EQ(EINVAL, kthread_stop(th));
        }
    }

    TEST(ButexTest, stop_before_sleeping) {
        turbo::TimeCost tm;
        const long SLEEP_MSEC = 100;

        for (int i = 0; i < 2; ++i) {
            kthread_t th;
            const kthread_attr_t attr =
                    (i == 0 ? KTHREAD_ATTR_PTHREAD : KTHREAD_ATTR_NORMAL) | KTHREAD_NOSIGNAL;

            tm.reset();
            ASSERT_EQ(0, kthread_start_background(&th, &attr, sleeper,
                                                  (void *) (SLEEP_MSEC * 1000L)));
            ASSERT_EQ(0, kthread_stop(th));
            kthread_flush();
            ASSERT_EQ(0, kthread_join(th, nullptr));
            tm.stop();

            if (attr.stack_type == KTHREAD_STACKTYPE_PTHREAD) {
                ASSERT_LT(labs(tm.m_elapsed() - SLEEP_MSEC), 10);
            } else {
                ASSERT_LT(tm.m_elapsed(), 10);
            }
            // ASSERT_TRUE(kthread::get_task_control()->
            //             timer_thread()._idset.empty());
            ASSERT_EQ(EINVAL, kthread_stop(th));
        }
    }

    void *trigger_signal(void *arg) {
        pthread_t *th = (pthread_t *) arg;
        const long t1 = turbo::Time::current_microseconds();
        for (size_t i = 0; i < 50; ++i) {
            usleep(100000);
            if (kthread::interrupt_pthread(*th) == ESRCH) {
                KLOG(INFO) << "waiter thread end, trigger count=" << i;
                break;
            }
        }
        const long t2 = turbo::Time::current_microseconds();
        KLOG(INFO) << "trigger signal thread end, elapsed=" << (t2 - t1) << "us";
        return nullptr;
    }

    TEST(ButexTest, wait_with_signal_triggered) {
        turbo::TimeCost tm;

        const int64_t WAIT_MSEC = 500;
        WaiterArg waiter_args;
        pthread_t waiter_th, tigger_th;
        std::atomic<int> *kutex =
                kthread::kutex_create_checked<std::atomic<int> >();
        ASSERT_TRUE(kutex);
        *kutex = 1;
        ASSERT_EQ(0, kthread::kutex_wake(kutex));

        const timespec abstime = turbo::Time::to_timespec(turbo::Time::future_time(turbo::Duration::milliseconds(WAIT_MSEC)));
        waiter_args.expected_value = *kutex;
        waiter_args.kutex = kutex;
        waiter_args.expected_result = ETIMEDOUT;
        waiter_args.ptimeout = &abstime;
        tm.reset();
        pthread_create(&waiter_th, nullptr, waiter, &waiter_args);
        pthread_create(&tigger_th, nullptr, trigger_signal, &waiter_th);

        ASSERT_EQ(0, pthread_join(waiter_th, nullptr));
        tm.stop();
        auto wait_elapsed_ms = tm.m_elapsed();;
        KLOG(INFO) << "waiter thread end, elapsed " << wait_elapsed_ms << " ms";

        ASSERT_LT(labs(wait_elapsed_ms - WAIT_MSEC), 250);

        ASSERT_EQ(0, pthread_join(tigger_th, nullptr));
        kthread::kutex_destroy(kutex);
    }

} // namespace
